import pybullet as p
import pybullet_data
import time
import numpy as np

# 连接到物理服务器
physicsClient = p.connect(p.GUI)
p.setGravity(0, 0, -9.8)

# 创建地面
p.createCollisionShape(p.GEOM_PLANE)
p.createMultiBody(0, 0)

# 添加数据路径
p.setAdditionalSearchPath(pybullet_data.getDataPath())

# 加载一个简单的机器人模型（比如KUKA机械臂）
robot_id = p.loadURDF("kuka_iiwa/model.urdf", [0, 0, 0])

# 获取关节数量和信息
num_joints = p.getNumJoints(robot_id)
print(f"机器人有 {num_joints} 个关节")

# 显示所有关节信息
for i in range(num_joints):
    joint_info = p.getJointInfo(robot_id, i)
    print(f"关节 {i}: {joint_info[1].decode('utf-8')}")

# 🎯 在 GUI 中添加调试参数控件
# 添加目标位置滑块
target_pos_slider = p.addUserDebugParameter(
    "目标位置",  # 参数名称（显示在UI中）
    -3.14,       # 最小值 (-π)
    3.14,        # 最大值 (π)
    0.0          # 初始值
)

# 添加控制力滑块
max_force_slider = p.addUserDebugParameter(
    "最大控制力", 
    0,           # 最小值
    500,         # 最大值
    200          # 初始值
)

# 添加控制增益滑块
position_gain_slider = p.addUserDebugParameter(
    "位置增益", 
    0.0,         # 最小值
    2.0,         # 最大值
    0.5          # 初始值
)

# 添加速度增益滑块
velocity_gain_slider = p.addUserDebugParameter(
    "速度增益", 
    0.0,         # 最小值
    2.0,         # 最大值
    1.0          # 初始值
)

# 添加关节选择器（选择控制哪个关节）
joint_selector = p.addUserDebugParameter(
    "控制关节", 
    0,           # 最小值（第一个关节）
    num_joints-1,# 最大值（最后一个关节）
    2            # 初始值（选择第三个关节）
)

# 添加控制模式切换按钮
control_mode_toggle = p.addUserDebugParameter(
    "控制模式 (0=位置, 1=速度)", 
    0,           # 最小值
    1,           # 最大值
    0            # 初始值（位置控制）
)

# 添加重置按钮
reset_button = p.addUserDebugParameter(
    "重置位置", 
    0,           # 最小值
    1,           # 最大值
    0            # 初始值
)

# 仿真循环
try:
    while True:
        # 🔍 读取所有调试参数的值
        target_position = p.readUserDebugParameter(target_pos_slider)
        max_force = p.readUserDebugParameter(max_force_slider)
        pos_gain = p.readUserDebugParameter(position_gain_slider)
        vel_gain = p.readUserDebugParameter(velocity_gain_slider)
        selected_joint = int(p.readUserDebugParameter(joint_selector))
        control_mode = int(p.readUserDebugParameter(control_mode_toggle))
        reset_value = p.readUserDebugParameter(reset_button)
        
        # 如果重置按钮被激活（值大于0.5），重置关节位置
        if reset_value > 0.5:
            for i in range(num_joints):
                p.resetJointState(robot_id, i, 0)
            # 将重置按钮的值设回0
            # 注意：PyBullet没有直接设置参数值的API，所以我们需要通过其他方式处理
        
        # 根据选择的控制模式设置关节控制
        if control_mode == 0:  # 位置控制
            p.setJointMotorControl2(
                bodyUniqueId=robot_id,
                jointIndex=selected_joint,
                controlMode=p.POSITION_CONTROL,
                targetPosition=target_position,
                force=max_force,
                positionGain=pos_gain,
                velocityGain=vel_gain
            )
        else:  # 速度控制
            p.setJointMotorControl2(
                bodyUniqueId=robot_id,
                jointIndex=selected_joint,
                controlMode=p.VELOCITY_CONTROL,
                targetVelocity=target_position,  # 这里用target_position作为目标速度
                force=max_force
            )
        
        # 推进仿真
        p.stepSimulation()
        
        # 获取当前关节状态并显示
        joint_state = p.getJointState(robot_id, selected_joint)
        joint_info = p.getJointInfo(robot_id, selected_joint)
        joint_name = joint_info[1].decode('utf-8')
        
        # 在3D空间中显示文本信息
        p.addUserDebugText(
            f"关节 {selected_joint} ({joint_name}):\n"
            f"位置: {joint_state[0]:.2f}\n"
            f"速度: {joint_state[1]:.2f}\n"
            f"力矩: {joint_state[3]:.2f}",
            textPosition=[0, 0, 1.5],
            textColorRGB=[1, 1, 0],
            textSize=1.2
        )
        
        time.sleep(1./240.)
        
except KeyboardInterrupt:
    pass

p.disconnect()